{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from pycocotools.coco import COCO\n",
    "from pycocotools.mask import iou, encode\n",
    "from pathlib import Path\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loading annotations into memory...\n",
      "Done (t=0.59s)\n",
      "creating index...\n",
      "index created!\n"
     ]
    }
   ],
   "source": [
    "prefix = 'instances'\n",
    "dataType = 'val2017'\n",
    "# Replace the following line with the actual location of your COCO dataset\n",
    "dataDir = Path('D:\\\\Documents\\\\AI\\\\COCO\\\\raw-data')\n",
    "\n",
    "annFile = dataDir/'annotations'/f'{prefix}_{dataType}.json'\n",
    "cocoGt = COCO(str(annFile)) # gt == \"ground truth\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision import transforms\n",
    "from PIL import Image\n",
    "\n",
    "preprocess = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "\n",
    "def get_image_filename(gt, imgId):\n",
    "    return str(dataDir/dataType/gt.imgs[imgId]['file_name'])\n",
    "\n",
    "def load_image(gt, imgId):\n",
    "    input_image = Image.open(get_image_filename(gt, imgId)).convert('RGB')\n",
    "    input_tensor = preprocess(input_image)\n",
    "    input_tensor = input_tensor.unsqueeze(0)\n",
    "    input_tensor = input_tensor.detach().cpu().numpy()\n",
    "    \n",
    "    return input_tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Dictionary to do map between different category ids\n",
    "# key = COCO id, value = Pascal VOC id\n",
    "COCO_TO_VOC = {\n",
    "    1: 15,  # person\n",
    "    2: 2,   # bicycle\n",
    "    3: 7,   # car\n",
    "    4: 14,  # motorbike\n",
    "    5: 1,   # airplane\n",
    "    6: 6,   # bus\n",
    "    7: 19,  # train\n",
    "    9: 4,   # boat\n",
    "    16: 3,  # bird\n",
    "    17: 8,  # cat\n",
    "    18: 12, # dog\n",
    "    19: 13, # horse\n",
    "    20: 17, # sheep\n",
    "    21: 10, # cow\n",
    "    44: 5,  # bottle\n",
    "    62: 9,  # chair\n",
    "    63: 18, # couch/sofa\n",
    "    64: 16, # potted plant\n",
    "    67: 11, # dining table\n",
    "    72: 20, # tv\n",
    "}\n",
    "\n",
    "VOC_CAT_IDS = list(COCO_TO_VOC.keys())\n",
    "\n",
    "def load_image_and_ann(gt, imgId):\n",
    "    input_tensor = load_image(gt, imgId)\n",
    "    \n",
    "    _, _, height, width = input_tensor.shape\n",
    "    output_tensor = np.zeros((21, height, width), dtype=np.uint8)\n",
    "    \n",
    "    annIds = gt.getAnnIds(imgId, VOC_CAT_IDS)\n",
    "    for ann in gt.loadAnns(annIds):\n",
    "        mask = gt.annToMask(ann)\n",
    "        output_tensor[COCO_TO_VOC[ann['category_id']]] |= mask\n",
    "        \n",
    "    # Set everything not labeled to be background\n",
    "    output_tensor[0] = 1 - np.max(output_tensor, axis=0)\n",
    "    \n",
    "    # Add extra dimension to comply with batch format\n",
    "    output_tensor = output_tensor[np.newaxis, ...]\n",
    "    \n",
    "    return input_tensor, output_tensor\n",
    "\n",
    "def getImgIdsUnion(gt, catIds):\n",
    "    \"\"\"\n",
    "    Returns all the images that have *any* of the categories in `catIds`,\n",
    "    unlike the built-in `gt.getImgIds` which returns all the images containing\n",
    "    *all* of the categories in `catIds`.\n",
    "    \"\"\"\n",
    "    imgIds = set()\n",
    "    \n",
    "    for catId in catIds:\n",
    "        imgIds |= set(gt.catToImgs[catId])\n",
    "        \n",
    "    return list(imgIds)\n",
    "\n",
    "imgIds = getImgIdsUnion(cocoGt, VOC_CAT_IDS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Completed 50/4031.\tmean IoU: 0.6345326703215275\tGPA: 0.9953693412234634\n",
      "Completed 100/4031.\tmean IoU: 0.650681260494401\tGPA: 0.9958247718348205\n",
      "Completed 150/4031.\tmean IoU: 0.6632394400667335\tGPA: 0.9956243610910763\n",
      "Completed 200/4031.\tmean IoU: 0.6635046302250455\tGPA: 0.9956633864336837\n",
      "Completed 250/4031.\tmean IoU: 0.6625868723429007\tGPA: 0.9957871990703088\n",
      "Completed 300/4031.\tmean IoU: 0.6548570232629858\tGPA: 0.9958517735336572\n",
      "Completed 350/4031.\tmean IoU: 0.657134582225294\tGPA: 0.9960188493331502\n",
      "Completed 400/4031.\tmean IoU: 0.6514943133576366\tGPA: 0.9956833655206868\n",
      "Completed 450/4031.\tmean IoU: 0.6551672826515782\tGPA: 0.9956050067966146\n",
      "Completed 500/4031.\tmean IoU: 0.6539223533792986\tGPA: 0.9956343288133468\n",
      "Completed 550/4031.\tmean IoU: 0.6513229173526096\tGPA: 0.9955664213994251\n",
      "Completed 600/4031.\tmean IoU: 0.6555730264528864\tGPA: 0.9956881602083523\n",
      "Completed 650/4031.\tmean IoU: 0.6588112248552683\tGPA: 0.9956474259746488\n",
      "Completed 700/4031.\tmean IoU: 0.6606245938352795\tGPA: 0.9956842402952785\n",
      "Completed 750/4031.\tmean IoU: 0.6614373338740641\tGPA: 0.9956420391729299\n",
      "Completed 800/4031.\tmean IoU: 0.6611653891164552\tGPA: 0.9957155101970937\n",
      "Completed 850/4031.\tmean IoU: 0.6626125995408813\tGPA: 0.9957210514017112\n",
      "Completed 900/4031.\tmean IoU: 0.663473155990766\tGPA: 0.9956990280281746\n",
      "Completed 950/4031.\tmean IoU: 0.663636521272714\tGPA: 0.9955901623197373\n",
      "Completed 1000/4031.\tmean IoU: 0.6643393869365718\tGPA: 0.9955910075507052\n",
      "Completed 1050/4031.\tmean IoU: 0.6635138441240718\tGPA: 0.9955672298917276\n",
      "Completed 1100/4031.\tmean IoU: 0.6621176945410658\tGPA: 0.9955764804700332\n",
      "Completed 1150/4031.\tmean IoU: 0.6631865107962432\tGPA: 0.9955986257651435\n",
      "Completed 1200/4031.\tmean IoU: 0.6657144101854745\tGPA: 0.99560734849275\n",
      "Completed 1250/4031.\tmean IoU: 0.6661664299148664\tGPA: 0.9956221203087662\n",
      "Completed 1300/4031.\tmean IoU: 0.6664217401048482\tGPA: 0.9956346196003822\n",
      "Completed 1350/4031.\tmean IoU: 0.6654635599026766\tGPA: 0.9956130023585732\n",
      "Completed 1400/4031.\tmean IoU: 0.6654526715951429\tGPA: 0.9956118096180874\n",
      "Completed 1450/4031.\tmean IoU: 0.6679818476008911\tGPA: 0.9956563423508025\n",
      "Completed 1500/4031.\tmean IoU: 0.6690841194049824\tGPA: 0.9956881763155492\n",
      "Completed 1550/4031.\tmean IoU: 0.6670264552011157\tGPA: 0.995667189162216\n",
      "Completed 1600/4031.\tmean IoU: 0.6665873625246779\tGPA: 0.9956424756554567\n",
      "Completed 1650/4031.\tmean IoU: 0.66671421902365\tGPA: 0.9956707905234181\n",
      "Completed 1700/4031.\tmean IoU: 0.6647246735470009\tGPA: 0.9956194850148153\n",
      "Completed 1750/4031.\tmean IoU: 0.6644462277095248\tGPA: 0.9956146820092521\n",
      "Completed 1800/4031.\tmean IoU: 0.6653214402942068\tGPA: 0.9956218160317079\n",
      "Completed 1850/4031.\tmean IoU: 0.666935355530244\tGPA: 0.9956436554687211\n",
      "Completed 1900/4031.\tmean IoU: 0.66671156788738\tGPA: 0.9956654951639361\n",
      "Completed 1950/4031.\tmean IoU: 0.668537352392893\tGPA: 0.995665754324596\n",
      "Completed 2000/4031.\tmean IoU: 0.6677127068842019\tGPA: 0.9956915301852337\n",
      "Completed 2050/4031.\tmean IoU: 0.6675186789692606\tGPA: 0.995676946325504\n",
      "Completed 2100/4031.\tmean IoU: 0.6675798019859625\tGPA: 0.9956999219785303\n",
      "Completed 2150/4031.\tmean IoU: 0.667148945007901\tGPA: 0.9956876831048003\n",
      "Completed 2200/4031.\tmean IoU: 0.6658898158479499\tGPA: 0.995662200864107\n",
      "Completed 2250/4031.\tmean IoU: 0.6644082238857807\tGPA: 0.9956967893537908\n",
      "Completed 2300/4031.\tmean IoU: 0.6656879132967034\tGPA: 0.9957077429742456\n",
      "Completed 2350/4031.\tmean IoU: 0.6655999286101026\tGPA: 0.9957157158290066\n",
      "Completed 2400/4031.\tmean IoU: 0.6663473375243069\tGPA: 0.9957214490273103\n",
      "Completed 2450/4031.\tmean IoU: 0.6668622259108372\tGPA: 0.9957133150862848\n",
      "Completed 2500/4031.\tmean IoU: 0.6665806350239287\tGPA: 0.9956962775891182\n",
      "Completed 2550/4031.\tmean IoU: 0.6666359274697989\tGPA: 0.9956701555589623\n",
      "Completed 2600/4031.\tmean IoU: 0.6665652228199483\tGPA: 0.9956606342709177\n",
      "Completed 2650/4031.\tmean IoU: 0.6661728124421915\tGPA: 0.9956544250763035\n",
      "Completed 2700/4031.\tmean IoU: 0.6655026100268836\tGPA: 0.9956311168798255\n",
      "Completed 2750/4031.\tmean IoU: 0.6653864299684282\tGPA: 0.9956338863411848\n",
      "Completed 2800/4031.\tmean IoU: 0.6662757800533626\tGPA: 0.9956426801150963\n",
      "Completed 2850/4031.\tmean IoU: 0.6646241929407267\tGPA: 0.9956475627505863\n",
      "Completed 2900/4031.\tmean IoU: 0.6654519349552466\tGPA: 0.9956739234829219\n",
      "Completed 2950/4031.\tmean IoU: 0.6658588186188478\tGPA: 0.995667829992838\n",
      "Completed 3000/4031.\tmean IoU: 0.6652108184297907\tGPA: 0.9956828292300053\n",
      "Completed 3050/4031.\tmean IoU: 0.6652147809862605\tGPA: 0.9956585649531894\n",
      "Completed 3100/4031.\tmean IoU: 0.6658580584456858\tGPA: 0.9956637690543105\n",
      "Completed 3150/4031.\tmean IoU: 0.6668207376988708\tGPA: 0.9956482604691415\n",
      "Completed 3200/4031.\tmean IoU: 0.6669557876916744\tGPA: 0.995640395956942\n",
      "Completed 3250/4031.\tmean IoU: 0.6662371236656779\tGPA: 0.995632590309063\n",
      "Completed 3300/4031.\tmean IoU: 0.6663868480407205\tGPA: 0.995629465369235\n",
      "Completed 3350/4031.\tmean IoU: 0.6667489868690694\tGPA: 0.9956187553104625\n",
      "Completed 3400/4031.\tmean IoU: 0.6672595782598338\tGPA: 0.9956403784125208\n",
      "Completed 3450/4031.\tmean IoU: 0.6681329206839\tGPA: 0.9956374377476654\n",
      "Completed 3500/4031.\tmean IoU: 0.6688084060506995\tGPA: 0.9956466605029343\n",
      "Completed 3550/4031.\tmean IoU: 0.6687079602696968\tGPA: 0.9956294135147684\n",
      "Completed 3600/4031.\tmean IoU: 0.6685197874758111\tGPA: 0.9956175511580541\n",
      "Completed 3650/4031.\tmean IoU: 0.6680068746548342\tGPA: 0.9956204426275244\n",
      "Completed 3700/4031.\tmean IoU: 0.6676773650423974\tGPA: 0.9956398778611952\n",
      "Completed 3750/4031.\tmean IoU: 0.6679168657352365\tGPA: 0.9956566207270627\n",
      "Completed 3800/4031.\tmean IoU: 0.6667399598858494\tGPA: 0.9956294387730882\n",
      "Completed 3850/4031.\tmean IoU: 0.6661352288486387\tGPA: 0.9956140241314541\n",
      "Completed 3900/4031.\tmean IoU: 0.6665341901309434\tGPA: 0.995608030477591\n",
      "Completed 3950/4031.\tmean IoU: 0.666812949939603\tGPA: 0.9956078423886188\n",
      "Completed 4000/4031.\tmean IoU: 0.6672221942307516\tGPA: 0.9956191774797091\n"
     ]
    }
   ],
   "source": [
    "from onnx import numpy_helper\n",
    "import os\n",
    "import onnxruntime as rt\n",
    "\n",
    "\n",
    "# Start from ORT 1.10, ORT requires explicitly setting the providers parameter if you want to use execution providers\n",
    "# other than the default CPU provider (as opposed to the previous behavior of providers getting set/registered by default\n",
    "# based on the build flags) when instantiating InferenceSession.\n",
    "# For example, if NVIDIA GPU is available and ORT Python package is built with CUDA, then call API as following:\n",
    "# rt.InferenceSession(path/to/model, providers=['CUDAExecutionProvider'])\n",
    "sess = rt.InferenceSession(\"../model/fcn-resnet101-11.onnx\")\n",
    "\n",
    "def predict(sess, input_tensor):\n",
    "    \"\"\"\n",
    "    Given an input tensor, create a (N, 21, height, width) one-hot\n",
    "    binary mask of pixelwise class predictions\n",
    "    \"\"\"\n",
    "    # The names for the FCN inputs/outputs are known, use them directly\n",
    "    model_tensor = sess.run([\"out\"], {\"input\": input_tensor})[0]\n",
    "    batch_size, nclasses, height, width = model_tensor.shape\n",
    "    \n",
    "    raw_labels = np.argmax(model_tensor, axis=1).astype(np.uint8)\n",
    "    # We need to convert from the argmax of each pixel into a one-hot binary tensor,\n",
    "    # which can be done with numpy's excellent boolean indexing\n",
    "    output_tensor = np.zeros((nclasses, batch_size, height, width), dtype=np.uint8)\n",
    "    \n",
    "    for c in range(nclasses):\n",
    "        output_tensor[c][raw_labels==c] = 1\n",
    "\n",
    "    output_tensor = np.transpose(output_tensor, [1, 0, 2, 3])\n",
    "    \n",
    "    return output_tensor\n",
    "\n",
    "def iou(model_tensor, target_tensor):\n",
    "    # Don't include the background when summing\n",
    "    model_tensor = model_tensor[:, 1:, :, :]\n",
    "    target_tensor = target_tensor[:, 1:, :, :]\n",
    "    \n",
    "    intersection = np.sum(np.logical_and(model_tensor, target_tensor))\n",
    "    union = np.sum(np.logical_or(model_tensor, target_tensor))\n",
    "    \n",
    "    if union == 0:\n",
    "        # Can only happen if nothing was there and nothing was predicted,\n",
    "        # which is a perfect score\n",
    "        return 1\n",
    "    else:\n",
    "        return intersection / union\n",
    "\n",
    "def pixelwise_accuracy(model_tensor, target_tensor):\n",
    "    batch_size, nclasses, height, width = model_tensor.shape\n",
    "    # Again, don't include the background in the accuracy\n",
    "    total_pixels = batch_size * (nclasses - 1) * height * width\n",
    "    model_tensor = model_tensor[:, 1:, :, :]\n",
    "    target_tensor = target_tensor[:, 1:, :, :]\n",
    "    \n",
    "    # Accuracy = (TP + TN) / (TP + TN + FP + FN)\n",
    "    return np.sum(model_tensor == target_tensor) / total_pixels\n",
    "\n",
    "def compute_validation_accuracy(gt, sess, imgIds, print_every=50):\n",
    "    \"\"\"\n",
    "    Given the ground truth annotations, compute the mean IoU and\n",
    "    global pixelwise accuracy for a given model instantiated in sess\n",
    "    on the image ids.\n",
    "    \n",
    "    This code is serial (non-batched) so it will be fairly slow.\n",
    "    \"\"\"\n",
    "    \n",
    "    totalIoU = 0\n",
    "    totalAcc = 0\n",
    "    totalImgs = len(imgIds)\n",
    "    \n",
    "    for i in range(totalImgs):\n",
    "        imgId = imgIds[i]\n",
    "        input_tensor, target_tensor = load_image_and_ann(gt, imgId)\n",
    "        model_tensor = predict(sess, input_tensor)\n",
    "\n",
    "        totalIoU += iou(model_tensor, target_tensor)\n",
    "        totalAcc += pixelwise_accuracy(model_tensor, target_tensor)\n",
    "        if print_every is not None and i % print_every == print_every - 1:\n",
    "            print(f\"Completed {i+1}/{totalImgs}.\\tmean IoU: {totalIoU / (i+1)}\\tGPA: {totalAcc / (i+1)}\")\n",
    "        \n",
    "    return totalIoU / totalImgs, totalAcc / totalImgs\n",
    "\n",
    "meanIoU, gpa = compute_validation_accuracy(cocoGt, sess, imgIds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "meanIoU=0.6673429850685226 gpa=0.9956271540146402\n"
     ]
    }
   ],
   "source": [
    "print(f\"{meanIoU=} {gpa=}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
